extern "C"
{
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
}

#include <opencv2/opencv.hpp>

#include <stdio.h>
#include <iostream>


void save_frame_as_image(AVFrame *frame, int width, int height, int frame_number) {
    char filename[256];
    snprintf(filename, sizeof(filename), "frame_%04d.jpg", frame_number);

    // 将AVFrame转换为OpenCV的Mat
    cv::Mat img(height, width, CV_8UC3, frame->data[0], frame->linesize[0]);

    // 保存为JPEG图像
    cv::imwrite(filename, img);
}

int main(void) {
    const char *input_filename = "test.mp4";

    // 打开输入文件
    AVFormatContext *fmt_ctx = NULL;
    if (avformat_open_input(&fmt_ctx, input_filename, NULL, NULL) != 0) {
        fprintf(stderr, "无法打开输入文件 %s\n", input_filename);
        return -1;
    }

    // 检索流信息
    if (avformat_find_stream_info(fmt_ctx, NULL) < 0) {
        fprintf(stderr, "无法检索流信息\n");
        return -1;
    }

    // 查找视频流
    int video_stream_index = -1;
    for (int i = 0; i < (int)fmt_ctx->nb_streams; i++) {
        if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_index = i;
            break;
        }
    }

    if (video_stream_index == -1) {
        fprintf(stderr, "未找到视频流\n");
        return -1;
    }

    // 获取视频流的编解码器参数
    AVCodecParameters *codec_params = fmt_ctx->streams[video_stream_index]->codecpar;

    // 找到解码器
    const AVCodec *decoder = avcodec_find_decoder(codec_params->codec_id);
    if (!decoder) {
        fprintf(stderr, "未找到解码器\n");
        return -1;
    }

    // 创建解码器上下文并打开解码器
    AVCodecContext *dec_ctx = avcodec_alloc_context3(decoder);
    if (avcodec_parameters_to_context(dec_ctx, codec_params) < 0) {
        fprintf(stderr, "无法复制解码器参数\n");
        return -1;
    }
    if (avcodec_open2(dec_ctx, decoder, NULL) < 0) {
        fprintf(stderr, "无法打开解码器\n");
        return -1;
    }

    // 分配AVFrame结构体
    AVFrame *frame = av_frame_alloc();
    AVFrame *rgb_frame = av_frame_alloc();
    if (!frame || !rgb_frame) {
        fprintf(stderr, "无法分配AVFrame\n");
        return -1;
    }

    // 设置SWScale上下文
    struct SwsContext *sws_ctx = sws_getContext(dec_ctx->width, dec_ctx->height, dec_ctx->pix_fmt,
                                                dec_ctx->width, dec_ctx->height, AV_PIX_FMT_BGR24,
                                                SWS_BILINEAR, NULL, NULL, NULL);

    // 开始解码循环
    int frame_count = 0;
    while (1) {
        AVPacket pkt;
        if (av_read_frame(fmt_ctx, &pkt) < 0)
            break;

        if (pkt.stream_index == video_stream_index) {
            int ret = avcodec_send_packet(dec_ctx, &pkt);
            if (ret < 0) {
                fprintf(stderr, "错误发送数据包给解码器\n");
                return -1;
            }

            while (ret >= 0) {
                ret = avcodec_receive_frame(dec_ctx, frame);
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                    break;
                else if (ret < 0) {
                    fprintf(stderr, "接收帧时出错\n");
                    return -1;
                }

                // 转换颜色空间
                av_image_alloc(rgb_frame->data, rgb_frame->linesize, dec_ctx->width, dec_ctx->height, AV_PIX_FMT_BGR24, 32);
                sws_scale(sws_ctx, (const uint8_t * const *)frame->data, frame->linesize, 0,
                          dec_ctx->height, rgb_frame->data, rgb_frame->linesize);

                // 保存为图片
                save_frame_as_image(rgb_frame, dec_ctx->width, dec_ctx->height, frame_count);

                av_freep(&rgb_frame->data[0]);
                frame_count++;
            }
        }
        av_packet_unref(&pkt);
    }

    // 清理
    av_frame_free(&frame);
    av_frame_free(&rgb_frame);
    avcodec_free_context(&dec_ctx);
    avformat_close_input(&fmt_ctx);
    sws_freeContext(sws_ctx);

    return 0;
}